#!/usr/bin/env python3
'''
遗传算法里面没有传统的什么通过梯度对网络参数进行调整
而是直接通过随机的方式对网络参数进行调整，然后通过评估
每次评估仅保留最好的网络
通过不断的迭代，最终得到可以达到目标奖励的网络
'''
import gymnasium as gym
import copy
import numpy as np

import torch
import torch.nn as nn

from tensorboardX import SummaryWriter

# 噪音影响网络参数的权重大小
NOISE_STD = 0.01
# 种群大小
POPULATION_SIZE = 50
PARENTS_COUNT = 10


class Net(nn.Module):
    def __init__(self, obs_size, action_size):
        super(Net, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_size, 32),
            nn.ReLU(),
            nn.Linear(32, action_size),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        return self.net(x)


def evaluate(env, net):
    '''
    评估种群

    返回奖励
    '''

    obs = env.reset()
    reward = 0.0
    while True:
        obs_v = torch.FloatTensor([obs])
        # 得到执行动作的预测
        act_prob = net(obs_v)
        acts = act_prob.max(dim=1)[1]
        # 执行动作得到奖励
        obs, r, done, _ = env.step(acts.data.numpy()[0])
        reward += r
        if done:
            break
    return reward


def mutate_parent(net):
    '''
    开始变异父母网络，通过将网络的参数加上噪声模拟父母变异的过程

    返回变异后（参数加入噪音）的网络
    '''
    #拷贝一个网络
    new_net = copy.deepcopy(net)
    for p in new_net.parameters():
        # 得到和网络参数相同维度的噪声
        noise_t = torch.tensor(np.random.normal(size=p.data.size()).astype(np.float32))
        # 将噪音添加到网络参数上
        p.data += NOISE_STD * noise_t
    return new_net


if __name__ == "__main__":
    writer = SummaryWriter(comment="-cartpole-ga")
    env = gym.make("CartPole-v0")

    gen_idx = 0
    # 创建对应种群大小的网络
    nets = [
        Net(env.observation_space.shape[0], env.action_space.n)
        for _ in range(POPULATION_SIZE)
    ]
    # 计算种群的奖励，population单词是人口，所以可能是
    # 通过奖励模拟人口的生存能力
    population = [
        (net, evaluate(env, net))
        for net in nets
    ]
    while True:
        # 得到降序排序的奖励（人口）
        population.sort(key=lambda p: p[1], reverse=True)
        # 得到前PARENTS_COUNT个奖励（人口）
        rewards = [p[1] for p in population[:PARENTS_COUNT]]
        # 分别计算平均值、最大值、标准差
        reward_mean = np.mean(rewards)
        reward_max = np.max(rewards)
        reward_std = np.std(rewards)

        writer.add_scalar("reward_mean", reward_mean, gen_idx)
        writer.add_scalar("reward_std", reward_std, gen_idx)
        writer.add_scalar("reward_max", reward_max, gen_idx)
        print("%d: reward_mean=%.2f, reward_max=%.2f, reward_std=%.2f" % (
            gen_idx, reward_mean, reward_max, reward_std))
        # 如果奖励的平均值大于199，就认为解决了问题
        if reward_mean > 199:
            print("Solved in %d steps" % gen_idx)
            break

        # generate next population
        # 预留奖励网络的顺序
        prev_population = population
        # 保留最好的一个网络
        population = [population[0]]
        # 循环计算遗传，循环的长度POPULATION_SIZE
        for _ in range(POPULATION_SIZE-1):
            # 得到随机的索引，索引范围是0到PARENTS_COUNT
            # parent_idx是父母的索引
            parent_idx = np.random.randint(0, PARENTS_COUNT)
            # 按照parent_idx得到父母的网络
            parent = prev_population[parent_idx][0]
            # 变异父母？将噪音添加进网络
            net = mutate_parent(parent)
            # 评估变异后的网络
            fitness = evaluate(env, net)
            # 将编译后的网络加入到种群中
            population.append((net, fitness))
        gen_idx += 1

    pass
